import os
import numpy as np
import mindspore
import mindspore.nn as nn
import mindspore.ops as ops
from mindspore import context, Tensor, dataset as ds
from mindspore.dataset.transforms import c_transforms
from mindspore.dataset.vision import c_transforms as vision
from mindspore.train import Model
from mindspore.train.callback import LossMonitor
from mindspore.nn import MSELoss, Adam

# 设置运行环境为 GPU
context.set_context(mode=context.GRAPH_MODE, device_target="GPU")


# 数据集加载与预处理
def create_image_dataset(image_folder, batch_size=10, image_size=(256, 256)):
    transforms = [
        vision.Decode(),
        vision.Resize(image_size),
        vision.Rescale(1.0 / 255.0, 0.0),
        vision.HWC2CHW()  # 转换为 (channels, height, width) 格式
    ]
    transform = c_transforms.Compose(transforms)
    dataset = ds.ImageFolderDataset(image_folder, shuffle=True)
    dataset = dataset.map(operations=transform, input_columns=["image"])
    dataset = dataset.batch(batch_size, drop_remainder=True)
    return dataset


# 定义秘密信息生成函数
def create_secret_dataset(num_samples, secret_size, batch_size):
    secrets = np.random.rand(num_samples, secret_size).astype(np.float32)
    secret_dataset = ds.NumpySlicesDataset({"secret": secrets}, shuffle=True)
    secret_dataset = secret_dataset.batch(batch_size, drop_remainder=True)
    return secret_dataset


# 失真层定义
class DistortionLayer(nn.Cell):
    def __init__(self):
        super(DistortionLayer, self).__init__()
        self.resize = ops.ResizeBilinear((300, 300))  # 缩放
        self.gaussian_noise = ops.StandardNormal()  # 高斯噪声
        self.random_crop = vision.RandomCrop((256, 256))  # 随机裁剪
        self.jpeg_op = vision.RandomSharpness()  # JPEG 压缩或图像模糊

    def construct(self, x):
        x = self.resize(x)  # 缩放图像
        x = x + self.gaussian_noise(x.shape) * 0.01  # 添加高斯噪声
        x = self.random_crop(x)  # 随机裁剪
        x = self.jpeg_op(x)  # 随机JPEG压缩或锐化
        return x


# StegaStampEncoder模型定义
class StegaStampEncoder(nn.Cell):
    def __init__(self):
        super(StegaStampEncoder, self).__init__()
        self.secret_dense = nn.Dense(100, 7500, activation=nn.ReLU(), weight_init='HeNormal')

        self.conv1 = nn.Conv2d(6, 32, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.conv2 = nn.Conv2d(32, 32, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.conv3 = nn.Conv2d(32, 64, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.conv4 = nn.Conv2d(64, 128, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.conv5 = nn.Conv2d(128, 256, kernel_size=3, stride=2, pad_mode='pad', padding=1)

        self.up6 = nn.Conv2d(256, 128, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.conv6 = nn.Conv2d(256, 128, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.up7 = nn.Conv2d(128, 64, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.conv7 = nn.Conv2d(128, 64, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.up8 = nn.Conv2d(64, 32, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.conv8 = nn.Conv2d(64, 32, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.up9 = nn.Conv2d(32, 32, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.conv9 = nn.Conv2d(70, 32, kernel_size=3, stride=1, pad_mode='pad', padding=1)
        self.residual = nn.Conv2d(32, 3, kernel_size=1, stride=1, pad_mode='pad', padding=0)

    def construct(self, inputs):
        secret, image = inputs
        secret = secret - 0.5
        image = image - 0.5

        secret = self.secret_dense(secret)
        secret = secret.view(secret.shape[0], 3, 50, 50)
        secret_enlarged = ops.ResizeBilinear((400, 400))(secret)

        inputs = ops.concat((secret_enlarged, image), axis=1)
        conv1 = ops.ReLU()(self.conv1(inputs))
        conv2 = ops.ReLU()(self.conv2(conv1))
        conv3 = ops.ReLU()(self.conv3(conv2))
        conv4 = ops.ReLU()(self.conv4(conv3))
        conv5 = ops.ReLU()(self.conv5(conv4))

        up6 = ops.ReLU()(self.up6(ops.ResizeBilinear((conv5.shape[2] * 2, conv5.shape[3] * 2))(conv5)))
        merge6 = ops.concat((conv4, up6), axis=1)
        conv6 = ops.ReLU()(self.conv6(merge6))

        up7 = ops.ReLU()(self.up7(ops.ResizeBilinear((conv6.shape[2] * 2, conv6.shape[3] * 2))(conv6)))
        merge7 = ops.concat((conv3, up7), axis=1)
        conv7 = ops.ReLU()(self.conv7(merge7))

        up8 = ops.ReLU()(self.up8(ops.ResizeBilinear((conv7.shape[2] * 2, conv7.shape[3] * 2))(conv7)))
        merge8 = ops.concat((conv2, up8), axis=1)
        conv8 = ops.ReLU()(self.conv8(merge8))

        up9 = ops.ReLU()(self.up9(ops.ResizeBilinear((conv8.shape[2] * 2, conv8.shape[3] * 2))(conv8)))
        merge9 = ops.concat((conv1, up9, inputs), axis=1)
        conv9 = ops.ReLU()(self.conv9(merge9))

        residual = self.residual(conv9)
        return residual


# WatermarkDecoder模型定义
class WatermarkDecoder(nn.Cell):
    def __init__(self, secret_size=100):
        super(WatermarkDecoder, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, pad_mode='pad', padding=1)
        self.fc1 = nn.Dense(256 * 32 * 32, 512)
        self.fc2 = nn.Dense(512, secret_size)

    def construct(self, x):
        x = ops.ReLU()(self.conv1(x))
        x = ops.ReLU()(self.conv2(x))
        x = ops.ReLU()(self.conv3(x))
        x = x.view(x.shape[0], -1)
        x = ops.ReLU()(self.fc1(x))
        x = self.fc2(x)
        return x


# 构建数据集
image_folder = "your_image_folder"  # 替换为你的图像文件夹路径
batch_size = 10
secret_size = 100

image_dataset = create_image_dataset(image_folder, batch_size=batch_size)
secret_dataset = create_secret_dataset(image_dataset.get_dataset_size() * batch_size, secret_size, batch_size)

# 定义损失函数和优化器
encoder = StegaStampEncoder()
decoder = WatermarkDecoder(secret_size=secret_size)
distortion = DistortionLayer()  # 失真层

loss_fn = MSELoss()
optimizer = Adam(encoder.trainable_params() + decoder.trainable_params(), learning_rate=0.001)


# 构建模型
class StegaStampNetwork(nn.Cell):
    def __init__(self, encoder, decoder, distortion):
        super(StegaStampNetwork, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.distortion = distortion

    def construct(self, image, secret):
        encoded_image = self.encoder((secret, image))
        distorted_image = self.distortion(encoded_image)  # 添加失真层
        decoded_secret = self.decoder(distorted_image)
        return decoded_secret


network = StegaStampNetwork(encoder, decoder, distortion)
model = Model(network, loss_fn, optimizer, metrics={"loss": nn.MSELoss()})

# 训练
epochs = 10
for epoch in range(epochs):
    for image_data, secret_data in zip(image_dataset.create_dict_iterator(), secret_dataset.create_dict_iterator()):
        image = Tensor(image_data['image'], mindspore.float32)
        secret = Tensor(secret_data['secret'], mindspore.float32)

        decoded_secret = model(image, secret)
        loss = loss_fn(decoded_secret, secret)

        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {loss.asnumpy()}")

print("训练结束！")
